from pathlib import Path

from loguru import logger

# from dialoggen.dialoggen_demo import DialogGen
from hydit.config import get_args
from hydit.inference_controlnet import End2End

from torchvision import transforms as T
import numpy as np
import os
import os.path as osp
import json
import torch
from tqdm import tqdm

# norm_transform = T.Compose(
#         [
#             T.ToTensor(),
#             T.Normalize([0.5], [0.5]),
#         ]
#     )

norm_transform = T.Compose(
        [
            T.ToTensor(),
            
        ]
    )

from PIL import Image
def inferencer():
    args = get_args()
    models_root_path = Path(args.model_root)
    if not models_root_path.exists():
        raise ValueError(f"`models_root` not exists: {models_root_path}")

    # Load models
    gen = End2End(args, models_root_path)

    # Try to enhance prompt
    if args.enhance:
        logger.info("Loading DialogGen model (for prompt enhancement)...")
        enhancer = DialogGen(str(models_root_path / "dialoggen"), args.load_4bit)
        logger.info("DialogGen model loaded.")
    else:
        enhancer = None

    return args, gen, enhancer


if __name__ == "__main__":

    args, gen, enhancer = inferencer()

    if enhancer:
        logger.info("Prompt Enhancement...")
        success, enhanced_prompt = enhancer(args.prompt)
        if not success:
            logger.info("Sorry, the prompt is not compliant, refuse to draw.")
            exit()
        logger.info(f"Enhanced prompt: {enhanced_prompt}")
    else:
        enhanced_prompt = None

    # Run inference
    logger.info("Generating images...")
    height, width = args.image_size

    # read tesdata
    promptdir = args.promptdir
    data = open(promptdir, 'r', encoding='utf-8')
    datalist = []
    for line in data.readlines():
        dic = json.loads(line)
        prompt = dic['text_zh']
        # if prompt.startswith('留白海报'):
        #         prompt = prompt.replace("留白海报", "海报背景")
        prompt = prompt[0:256]
        imgpath = dic['imgpath']
        maskpath = dic['maskpath']
        datalist.append((imgpath, maskpath, prompt))
    print('test imgs', len(datalist))


    savedir = args.savedir
    os.makedirs(savedir, exist_ok=True)


    # maskdir = 'mask.jpg'
    # mask = (np.array(Image.open(maskdir)))/255
    # mask = mask[:, :,np.newaxis]
    # mask[mask < 0.5] = 0
    # mask[mask >= 0.5] = 1
    # mask = mask.transpose(2, 0, 1) # c,h,w
    # mask = torch.from_numpy(mask)  # tensor:1,h,w, [0, 1]

    # condition = Image.open(args.condition_image_path).convert('RGB').resize((height, width))
    # image = norm_transform(condition)

    # masked = (1-mask)*image
    # masked = masked.unsqueeze(0).cuda()

    for item in tqdm(datalist):
        try:
            imgpath, maskpath, caption = item
            mask = Image.open(maskpath)
            mask = mask.resize((1024,1024))
            mask = np.array(mask)/255
            # mask = mask[:, :,np.newaxis]
            mask[mask < 0.5] = 0
            mask[mask >= 0.5] = 1
            mask_np = mask
            mask = mask.transpose(2, 0, 1) # c,h,w
            mask = torch.from_numpy(mask)  # tensor:1,h,w, [0, 1]

            masked = Image.open(imgpath).convert('RGB')
            masked = masked.resize((1024, 1024))
            masked_np = np.array(masked)
            w,h = masked.size
            condition = norm_transform(masked)

            condition = (1-mask)*condition
            condition = condition * 2.0 - 1.0
            condition = condition.unsqueeze(0).cuda()

            # 艺术海报,高质量
            # "photograph of a beautiful empty scene, highest quality settings"
            # "艺术海报,高质量"
            prompt = '艺术海报,高质量'
            # prompt = caption
            for i in range(1):
        
                results = gen.predict(prompt,
                                    height=h,
                                    width=w,
                                    image=condition,
                                    seed=None,
                                    enhanced_prompt=enhanced_prompt,
                                    negative_prompt=args.negative,
                                    infer_steps=args.infer_steps,
                                    guidance_scale=args.cfg_scale,
                                    batch_size=args.batch_size,
                                    src_size_cond=(w, h),
                                    use_style_cond=args.use_style_cond,
                                    style=args.style
                                    )['images'][0]
                if len(caption)>100:
                    saveprompt=caption[0:100]
                else:
                    saveprompt = caption
                output = (np.array(results)/255)*mask_np + (masked_np/255)*(1-mask_np)
                output = (output*255).astype(np.uint8)
                output = Image.fromarray(output)
                output.save(osp.join(savedir, saveprompt + '_' + str(i) + '.jpg'))
        except Exception as e:
            print(e)
            continue


